import pandas as pd

# 设置参数
N = 9
M1 = 3
M2 = 3


def test():
    pass


# 定义函数计算LLV（最低价最小值）和HHV（最高价最大值）
def calculate_llv_high_low(high, low, n):
    return low.rolling(window=n).min(), high.rolling(window=n).max()


# 定义函数计算RSV
def calculate_rsv(close, llv_low, hhv_high):
    return (close - llv_low) / (hhv_high - llv_low) * 100


# 定义函数计算SMA（简单移动平均）
def calculate_sma(data, window):
    return data.rolling(window=window).mean()


# 获取kdj信息
def get_kdj(df):
    # 计算LLV和HHV
    df['LLV_LOW'], df['HHV_HIGH'] = calculate_llv_high_low(df['HIGH'], df['LOW'], N)

    # 计算RSV
    df['RSV'] = calculate_rsv(df['CLOSE'], df['LLV_LOW'], df['HHV_HIGH'])

    # 计算K
    df['K'] = calculate_sma(df['RSV'], M1)

    # 计算D
    df['D'] = calculate_sma(df['K'], M2)

    # 计算J
    df['J'] = 3 * df['K'] - 2 * df['D']

    return df
